#!/usr/bin/python3
#-*-coding:utf8-*-

from transformers import pipeline
import web

pipe = pipeline('text-generation', model='daspartho/prompt-extend')
#pipe('cat,', num_return_sequences=1,max_length=10)[0]["generated_text"]

urls = (
    '/prompt', 'prompt'
)

app = web.application(urls, globals())

class prompt:
    def POST(self):
        p = web.input().prompt
        return pipe(p + ',', num_return_sequences=1)[0]["generated_text"]   
        #return pipe(p + ',', num_return_sequences=1,max_length=10)[0]["generated_text"]   


if __name__ == "__main__":
    app.run()  

